import requests
import json
import time
import os

# 配置
GRAPHQL_ENDPOINT = "http://192.168.49.2:30005/graphql"
COOKIES = {
    'JSESSIONID': '1BCCECD1F7162EFF2D592E169F199497',
    'grafana_session': '015c3c3e94d17323763cf638c25c2c3e',
}

HEADERS = {
    'Accept': 'application/json, text/plain, */*',
    'Content-Type': 'application/json;charset=UTF-8',
    'Origin': 'http://localhost:30005',
    'Referer': 'http://localhost:30005/trace',
    'User-Agent': 'Mozilla/5.0',
}




# 获取 traceIds
def get_trace_ids(START_TIME, END_TIME):
    """调用 queryBasicTraces 获取 traceIds 列表"""
    query = '''
    query queryTraces($condition: TraceQueryCondition) {
      data: queryBasicTraces(condition: $condition) {
        traces {
          traceIds
        }
      }
    }
    '''
    variables = {
        "condition": {
            "queryDuration": {
                "start": START_TIME,
                "end": END_TIME,
                "step": "MINUTE"
            },
            "traceState": "ALL",
            "paging": {
                "pageNum": 1,
                "pageSize": 9999,  # 可根据需要扩展
                "needTotal": False
            },
            "queryOrder": "BY_START_TIME"
        }
    }

    response = requests.post(GRAPHQL_ENDPOINT, cookies=COOKIES, headers=HEADERS, json={
        "query": query,
        "variables": variables
    })

    data = response.json()
    # print(json.dumps(data, indent=2, ensure_ascii=False))  # 打印原始数据
    traces = data["data"]["data"]["traces"]

    # traceId去重
    trace_ids_set = set()
    trace_ids = []

    for trace in traces:
        if trace["traceIds"]:
            trace_id = trace["traceIds"][0]
            if trace_id not in trace_ids_set:
                trace_ids.append(trace_id)
                trace_ids_set.add(trace_id)

    return trace_ids

# 获取 trace 详细信息
def get_trace_detail(trace_id):
    """调用 queryTrace 获取详细 spans"""
    query = '''
    query queryTrace($traceId: ID!) {
      trace: queryTrace(traceId: $traceId) {
        spans {
          traceId
          segmentId
          spanId
          parentSpanId
          refs {
            traceId
            parentSegmentId
            parentSpanId
            type
          }
          serviceCode
          serviceInstanceName
          startTime
          endTime
          endpointName
          type
          peer
          component
          isError
          layer
          tags {
            key
            value
          }
          logs {
            time
            data {
              key
              value
            }
          }
        }
      }
    }
    '''

    variables = {"traceId": trace_id}
    response = requests.post(GRAPHQL_ENDPOINT, cookies=COOKIES, headers=HEADERS, json={
        "query": query,
        "variables": variables
    })

    return response.json()["data"]["trace"]["spans"]

def format_trace(trace_id, spans):
    """格式化为目标 JSON 结构"""
    formatted_spans = []
    processes = {}

    for span in spans:
        formatted_spans.append({
            "traceID": span["traceId"],
            "spanID": str(span["spanId"]),
            "flags": 1,
            "operationName": span.get("endpointName", ""),
            "references": [
                {
                    "refType": ref["type"],
                    "traceID": ref["traceId"],
                    "spanID": str(ref["parentSpanId"]),
                } for ref in span.get("refs", [])
            ],
            "startTime": span["startTime"],
            "duration": span["endTime"] - span["startTime"],
            "tags": [
                {
                    "key": tag["key"],
                    "type": "string",
                    "value": tag["value"]
                } for tag in span.get("tags", [])
            ],
            "logs": [],
            "processID": "p1",
            "warnings": None
        })

        processes["p1"] = {
            "serviceName": span["serviceCode"],
            "tags": [
                {"key": "instance", "type": "string", "value": span["serviceInstanceName"]},
                {"key": "layer", "type": "string", "value": span["layer"]},
            ]
        }

    return {
        "traceID": trace_id,
        "spans": formatted_spans,
        "processes": processes,
        "warnings": None
    }

def get_spans(START_TIME, END_TIME, path):

    # 设置时间范围
    # START_TIME = "2025-05-22 0537"
    # END_TIME = "2025-05-22 0552"
    trace_ids = get_trace_ids(START_TIME=START_TIME, END_TIME=END_TIME)
    print(f"获取到 {len(trace_ids)} 条 trace")

    all_traces = []

    for trace_id in trace_ids:
        # print(f"处理 traceId: {trace_id}")
        spans = get_trace_detail(trace_id)
        formatted_trace = format_trace(trace_id, spans)
        all_traces.append(formatted_trace)
        time.sleep(0.2)  # 防止请求过快

    spans_path = os.path.join(path, "spans.json")

    with open(spans_path, "w", encoding="utf-8") as f:
        json.dump(all_traces, f, indent=2)

    print("已保存到 spans.json")


